#!/usr/bin/env python
# _*_ coding: utf-8 _*_
#
#
# @File : exdense.py.py
# @Time : 2020-09-17 23:03 
# Copyright (C) 2020 WeiKeting<weikting@gmail.com>. All rights reserved.
# @Description :
#
#

import tensorflow as tf
import tensorflow.keras.backend as K


class ExtendDense0(tf.keras.layers.Layer):

    def __init__(self,
                 units,
                 activation=None,
                 kernel_initializer='glorot_uniform',
                 kernel_regularizer=None,
                 activity_regularizer=None,
                 kernel_constraint=None,
                 normalize=True,
                 scale=16,
                 **kwargs):
        super().__init__(
            activity_regularizer=activity_regularizer, **kwargs)

        from tensorflow.keras import activations, initializers, regularizers, constraints
        from tensorflow.python.keras.engine.base_layer import InputSpec

        self.units = int(units) if not isinstance(units, int) else units

        self.activation = activations.get(activation)
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.kernel_regularizer = regularizers.get(kernel_regularizer)
        self.kernel_constraint = constraints.get(kernel_constraint)

        self.input_spec = InputSpec(min_ndim=2)
        self.supports_masking = True
        self._normalize = normalize
        self._scale = scale

    def build(self, input_shape):
        from tensorflow.python.framework import dtypes
        from tensorflow.python.framework import tensor_shape
        from tensorflow.python.layers.base import InputSpec

        dtype = dtypes.as_dtype(self.dtype or K.floatx())
        if not (dtype.is_floating or dtype.is_complex):
            raise TypeError('Unable to build `Dense` layer with non-floating point '
                            'dtype %s' % (dtype,))

        input_shape = tensor_shape.TensorShape(input_shape)
        last_dim = tensor_shape.dimension_value(input_shape[-1])
        if last_dim is None:
            raise ValueError('The last dimension of the inputs to `Dense` '
                             'should be defined. Found `None`.')
        self.input_spec = InputSpec(min_ndim=2, axes={-1: last_dim})
        self.kernel = self.add_weight(
            'kernel',
            shape=[last_dim, self.units],
            initializer=self.kernel_initializer,
            regularizer=self.kernel_regularizer,
            constraint=self.kernel_constraint,
            dtype=self.dtype,
            trainable=True)
        self.built = True

    def call(self, inputs, **kwargs):
        w = self.kernel
        x = inputs
        if self._normalize:
            w = K.l2_normalize(w, axis=0)
            x = K.l2_normalize(x, axis=1)
        return tf.matmul(x, w) * self._scale

    def compute_output_shape(self, input_shape):
        from tensorflow.python.framework import tensor_shape
        input_shape = tensor_shape.TensorShape(input_shape)
        input_shape = input_shape.with_rank_at_least(2)
        if tensor_shape.dimension_value(input_shape[-1]) is None:
            raise ValueError(
                'The innermost dimension of input_shape must be defined, but saw: %s'
                % input_shape)
        return input_shape[:-1].concatenate(self.units)

    def get_config(self):
        from tensorflow.keras import activations, initializers, regularizers, constraints
        config = super().get_config()
        config.update({
            'units':
                self.units,
            "normalize": self._normalize,
            "scale": self._scale,
            'activation':
                activations.serialize(self.activation),
            'kernel_initializer':
                initializers.serialize(self.kernel_initializer),
            'kernel_regularizer':
                regularizers.serialize(self.kernel_regularizer),
            'activity_regularizer':
                regularizers.serialize(self.activity_regularizer),
            'kernel_constraint':
                constraints.serialize(self.kernel_constraint),
        })
        return config


class ExtendDense1(tf.keras.layers.Dense):

    def __init__(self, units, normalize=True, scale=16, **kwargs):
        kwargs['use_bias'] = False
        super().__init__(units, **kwargs)
        self._normalize = normalize
        self._scale = scale

    def build(self, input_shape):
        super().build(input_shape)
        if self._normalize:
            self._kernel = self.kernel
            self.kernel = tf.keras.backend.l2_normalize(self._kernel) * self._scale

    def call(self, inputs):
        if self._normalize:
            inputs = tf.keras.backend.l2_normalize(inputs, axis=1)
        return super().call(inputs)


ExtendDense = ExtendDense0
